package phototagger.imgsim;

import java.lang.Math;
import java.lang.System;

/** Does the Haar-basis 2D wavelet decomposition of the image 
    (represented as a 2D array of doubles). This function is color-space
    agnostic, but in my implementation it is only used on the YIQ color
    space.
    
    @see "Wavelets for Computer Graphics", chapter 3
**/
public class HaarWavelet {

    public static boolean DEBUG = false;

    /** Performs the standard Haar basis wavelet decomposition.
	@param image the Image to decompose
	@returns an Image object with the columns decomposed
     **/
    public YIQImage decompose(YIQImage image) {

	String colors = image.getColorSpace();
	YIQImage decomposedImage = new YIQImage();

	for (int i = 0; i < colors.length(); i++) {
	    char color = colors.charAt(i);

	    double[][] decomposedElements =
		standardDecomposition(image.getElements(color));
	    
	    decomposedImage.setElements(color, decomposedElements);
	    decomposedImage.setAverage(color, decomposedElements[0][0]);
	}

	return decomposedImage;
    }
    
    /** I just KNOW this method is going to be riddled with bugs :( **/

    public double[][] standardDecomposition(double[][] colorComponent) 
	throws IllegalArgumentException {
	
	//check to see if the array has  2^j rows and 2^k columns, 
	// error out. Actually: 2^j -1, (0-127).
	if(colorComponent.length % 2 != 0) {
	    String m = "Rows not mod 2. Number of rows: " +
		colorComponent.length;
	    throw new IllegalArgumentException(m);
	} else if ((colorComponent[0]).length  % 2 != 0) {
	    String m = "Columns not mod 2. Number of columns is " +
		(colorComponent[0]).length;
	    throw new IllegalArgumentException(m);
	} else if(colorComponent.length != (colorComponent[0]).length) {
	    String m = "Number of rows and columns not equal.";
	    throw new IllegalArgumentException(m);
	}
	
	for(int row=0; row<colorComponent.length; row++) {
	    colorComponent[row] = decomposition(colorComponent[row]); 
	}
	
	// for each possible column element...
	for(int col=0; col<(colorComponent[0]).length; col++) {
	    
	    double[] column = new double[colorComponent.length];
	    
	    // for each row element, extract the col-th column element
	    for(int row=0; row<(colorComponent.length); row++) {
		column[row] = colorComponent[row][col];
	    }
	    
	    if(DEBUG) {
		for(int i=0; i<column.length ;i++) {
		    System.out.println("column["+i+"] = " + column[i]);
		}
	    }
	    
	    column = decomposition(column);
	    
	    // now, put that column back into colorComponent where 
	    // it belongs
	    for(int i=0; i<column.length; i++) {
		colorComponent[i][col] = column[i];
	    }
	}
	
	return colorComponent;
    }

    public double[] decomposition(double[] elements) throws
	IllegalArgumentException {
	// check to see that there are 2^j elements...and definately
	// not 0!!!
	if ((elements.length % 2 != 0) || elements.length == 0) {
	    String m = "There must be 2^j elements in decomposition";
	    throw new IllegalArgumentException(m);
	}
	
	// normalize input coeffients;
	double factor = Math.sqrt((double)elements.length);
	for(int i=0; i<elements.length; i++) {
	    elements[i] = elements[i] / factor;
	}
	
	int g = elements.length;
	
	while(g >= 2) {
	    double[] temp = new double[g];
	    System.arraycopy(elements, 0, temp, 0, g);
	    temp = decompositionStep(temp); 
	    System.arraycopy(temp, 0, elements, 0, temp.length);
	    g = g / 2;
	}

	return elements;
    }
    
    /** Watch carefully for off-by-one errors here... **/
    public double[] decompositionStep(double[] elements) {
	double root2 = Math.sqrt(2d);
	double[] newElements = new double[elements.length];
	for(int i=0; i<elements.length/2; i++) {
	    newElements[i] = (elements[2*i] + elements[2*i + 1]) / root2;
	    newElements[elements.length/2 + i] = 
		(elements[2*i] - elements[2*i + 1]) / root2;
	}
	    
	return newElements;
    }

}
